{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r9D9r4VxKe1-"
      },
      "source": [
        "Copyright 2021 Google LLC\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "    https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EzApNWZqKi-9"
      },
      "source": [
        "# Recursively Fertile Self-replicating Neural Agents\n",
        "\n",
        "Code base for reproducing the results and extra material of \"Recursively Fertile Self-replicating Neural Agents\", by Ettore Randazzo, Luca Versari and Alexander Mordvintsev. Paper published in ALIFE 2021."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8f0tihlDJe7K"
      },
      "outputs": [],
      "source": [
        "#@title Imports and Notebook Utilities\n",
        "%tensorflow_version 2.x\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pylab as pl\n",
        "import tensorflow as tf\n",
        "\n",
        "import io\n",
        "import math\n",
        "import pathlib\n",
        "import PIL\n",
        "from IPython.display import Image, clear_output\n",
        "\n",
        "\n",
        "def np2pil(a):\n",
        "  if a.dtype in [np.float32, np.float64]:\n",
        "    a = np.uint8(np.clip(a, 0, 1)*255)\n",
        "  return PIL.Image.fromarray(a)\n",
        "\n",
        "def imwrite(f, a, fmt=None):\n",
        "  a = np.asarray(a)\n",
        "  if isinstance(f, str):\n",
        "    fmt = f.rsplit('.', 1)[-1].lower()\n",
        "    if fmt == 'jpg':\n",
        "      fmt = 'jpeg'\n",
        "    f = open(f, 'wb')\n",
        "  np2pil(a).save(f, fmt, quality=95)\n",
        "\n",
        "def imencode(a, fmt='jpeg'):\n",
        "  a = np.asarray(a)\n",
        "  if len(a.shape) == 3 and a.shape[-1] == 4:\n",
        "    fmt = 'png'\n",
        "  f = io.BytesIO()\n",
        "  imwrite(f, a, fmt)\n",
        "  return f.getvalue()\n",
        "\n",
        "def imshow(a, fmt='jpeg'):\n",
        "  display(Image(data=imencode(a, fmt)))\n",
        "\n",
        "def zoom(img, scale=4):\n",
        "  img = np.repeat(img, scale, 0)\n",
        "  img = np.repeat(img, scale, 1)\n",
        "  return img\n",
        "\n",
        "\n",
        "!nvidia-smi -L"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mZCDkL1rLxgh"
      },
      "source": [
        "# Prepare the target image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c8lbUfbjMAKi"
      },
      "outputs": [],
      "source": [
        "# parameters used in the resizing.\n",
        "IMG_HEIGHT = 128\n",
        "IMG_WIDTH = IMG_HEIGHT\n",
        "\n",
        "def show_single_im(im, figsize=(3,3), reshape=False):\n",
        "  if reshape:\n",
        "    im = im.reshape([IMG_HEIGHT, IMG_WIDTH, 3])\n",
        "  pl.figure(figsize=figsize)\n",
        "  pl.imshow(im)\n",
        "  pl.axis(\"off\")\n",
        "  pl.show()\n",
        "\n",
        "\n",
        "def decode_img(img):\n",
        "  # convert the compressed string to a 3D uint8 tensor\n",
        "  img = tf.image.decode_jpeg(img, channels=3)\n",
        "  # Use `convert_image_dtype` to convert to floats in the [0,1] range.\n",
        "  img = tf.image.convert_image_dtype(img, tf.float32)\n",
        "  # resize the image to the desired size.\n",
        "  return tf.image.resize(img, [IMG_WIDTH, IMG_HEIGHT])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dopqc7siMJmw"
      },
      "outputs": [],
      "source": [
        "# We chose to use a rose as a target image. Here you can see more of them in\n",
        "# case you are interested in testing this work with other images.\n",
        "data_dir = tf.keras.utils.get_file(\n",
        "    origin='https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',\n",
        "    fname='flower_photos', untar=True)\n",
        "data_dir = pathlib.Path(data_dir)\n",
        "\n",
        "image_count = len(list(data_dir.glob('*/*.jpg')))\n",
        "print(\"Number of images in dataset:\", image_count)\n",
        "CLASS_NAMES = np.array([item.name for item in data_dir.glob('*') if item.name != \"LICENSE.txt\"])\n",
        "print(\"Classes:\", CLASS_NAMES)\n",
        "\n",
        "roses = list(sorted(data_dir.glob('roses/*')))\n",
        "print(\"Example roses:\")\n",
        "for image_path in roses[:10]:\n",
        "    imshow(PIL.Image.open(str(image_path)))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QekNpyq6NwcR"
      },
      "outputs": [],
      "source": [
        "# the target we choose to optimize for:\n",
        "target = tf.reshape(decode_img(tf.io.read_file(str(roses[9]))), [-1, 3])\n",
        "show_single_im(target.numpy(), figsize=(10, 10), reshape=True)\n",
        "\n",
        "# Also generate the input values for the image encoding task.\n",
        "r = 1.0 # change this if you want to have different ranges.\n",
        "coord_range = tf.linspace(-r, r, IMG_HEIGHT)\n",
        "y, x = tf.meshgrid(coord_range, coord_range, indexing='ij')\n",
        "im_enc_input = tf.reshape(tf.stack([x, y], -1), [-1, 2])\n",
        "print(\"target shape:\", target.shape)\n",
        "print(\"im_enc_input shape:\", im_enc_input.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HCSrhonm5zJu"
      },
      "outputs": [],
      "source": [
        "# @title plot utils\n",
        "\n",
        "def show_weight_divergence(all_w):\n",
        "  wd_fp_l, wd_fl_l = [], []\n",
        "  w_first, w_last = all_w[0], all_w[-1]\n",
        "  for idx, w in enumerate(all_w):\n",
        "    if idx \u003e 0:\n",
        "      wd_fp = tf.reduce_mean(tf.square(all_w[idx-1] - w))\n",
        "      wd_fp_l.append(wd_fp)\n",
        "    if idx \u003c len(all_w)-1:\n",
        "      wd_last = tf.reduce_mean(tf.square((w_last - w)))\n",
        "      wd_fl_l.append(wd_last)\n",
        "\n",
        "  pl.title(\"Weight divergence\")\n",
        "  x_l = np.arange(0, len(wd_fp_l)+1, dtype=np.int32)\n",
        "  x_l_1_to_n = x_l[1:]\n",
        "  x_l_0_to_nm1 = x_l[:-1]\n",
        "  pl.plot(x_l_1_to_n, wd_fp_l, label=\"From parent\")\n",
        "  pl.plot(x_l_0_to_nm1, wd_fl_l, label=\"From last (sink)\")\n",
        "  pl.yscale(\"log\")\n",
        "  pl.xlabel(\"Number of replication steps\")\n",
        "  pl.legend()\n",
        "  pl.show()\n",
        "\n",
        "  p_data = [(x_l_1_to_n, wd_fp_l), (x_l_0_to_nm1, wd_fl_l)]\n",
        "  return p_data\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qjBEC8ATcWhB"
      },
      "source": [
        "# Model architecture"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IbNE3g4CaB-k"
      },
      "outputs": [],
      "source": [
        "from tensorflow.keras.layers import Dense\n",
        "\n",
        "def getNumBits(n):\n",
        "  return int(math.floor(math.log(n, 2.0)) + 1)\n",
        "\n",
        "class SelfReplicator():\n",
        "  def __init__(self, n_hidden, wo, size_hidden, standardize_weights,\n",
        "               use_fixed_weights, switch_init_constant, num_extra_params):\n",
        "    \"\"\"\n",
        "    params:\n",
        "    n_hidden: number of hidden layers\n",
        "    wo: sinusoidal period for the first layer. Taken by \n",
        "      \"Implicit Neural Representations with Periodic Activation Functions\",\n",
        "      Sitzmann et al.\n",
        "    size_hidden: size of all hidden layers\n",
        "    standardize_weights: whether to standardize weights during self-replication.\n",
        "    use_fixed_weights: whether the network also uses fixed weights (exp 2 and 3)\n",
        "    switch_init_constant: value of the W_switch (here self.is_fixed_logit) \n",
        "      for initialization.\n",
        "    num_extra_params: used in exp3 to add replication with variation inputs.\n",
        "\n",
        "    \"\"\"\n",
        "    self.n_hidden = n_hidden\n",
        "    self.size_hidden = size_hidden\n",
        "    self.standardize_weights = standardize_weights\n",
        "    self.input_aux_size = 3\n",
        "    self.extra_params = num_extra_params\n",
        "    self.n_outputs = 4 # rgb + w\n",
        "    self.input_coord_size = self._computeInputCoordSize() \n",
        "    self.layers_sizes = [self.input_aux_size + self.input_coord_size + self.extra_params] + [self.size_hidden]*self.n_hidden + [self.n_outputs]\n",
        "    self.wo = wo\n",
        "    self.use_fixed_weights = use_fixed_weights\n",
        "    self.switch_init_constant = switch_init_constant\n",
        "    self.trainable_variables = []\n",
        "    layers = []\n",
        "    n = self.input_aux_size + self.input_coord_size + self.extra_params\n",
        "    for i in range(n_hidden):\n",
        "      n2 = size_hidden\n",
        "      k = math.sqrt(6/n)\n",
        "      layer_init = np.random.uniform(-k, k, [n, n2]).astype(np.float32)\n",
        "      if i == 0:\n",
        "        layer_init *= wo\n",
        "      layer_init = tf.Variable(layer_init)\n",
        "      self.trainable_variables.append(layer_init)\n",
        "      bias_init = tf.Variable(np.zeros(n2, dtype=np.float32))\n",
        "      self.trainable_variables.append(bias_init)\n",
        "      layers.append((layer_init, bias_init))\n",
        "      n = n2\n",
        "    \n",
        "    n2 = self.n_outputs\n",
        "    k = math.sqrt(6/n)\n",
        "    layer_init = np.random.uniform(-k, k, [n, n2]).astype(np.float32)\n",
        "    layer_init = tf.Variable(layer_init)\n",
        "    self.trainable_variables.append(layer_init)\n",
        "    layers.append((layer_init,))\n",
        "\n",
        "    self.layers = layers\n",
        "    self.layer_inputs = self._createBinaryCoords()\n",
        "    self.all_li = tf.concat(self.layer_inputs, 0)\n",
        "\n",
        "    if self.use_fixed_weights:\n",
        "      all_weights = tf.concat([tf.reshape(v, [-1]) for v in self.trainable_variables], 0)\n",
        "      switch_np = np.full(all_weights.shape, self.switch_init_constant).astype(np.float32)\n",
        "      self.is_fixed_logit = tf.Variable(switch_np)\n",
        "      self.fixed_w = tf.Variable(all_weights)\n",
        "      self.trainable_fixed_variables = [self.is_fixed_logit, self.fixed_w]\n",
        "\n",
        "  def _computeInputCoordSize(self):\n",
        "    # this is tricky because the first layer size depends on the input coord size.\n",
        "    # therefore, we have to find a stable configuration.\n",
        "    # We start from the second layers since these ones have fixed size.\n",
        "    ls_except_first = [self.size_hidden]*self.n_hidden + [self.n_outputs]\n",
        "    ics = 0\n",
        "    for i in range(len(ls_except_first) -1):\n",
        "      ls1, ls2 = ls_except_first[i], ls_except_first[i+1]\n",
        "      ics += getNumBits(ls1*ls2)\n",
        "      if i \u003c len(ls_except_first) - 2:\n",
        "        # has bias too\n",
        "        ics += getNumBits(ls2)\n",
        "    \n",
        "    # the first layer bias is also fixed in size:\n",
        "    ics += getNumBits(self.size_hidden)\n",
        "\n",
        "    # now see what is the minimum number of bits that fits the weight matrix.\n",
        "    bit_capacity = 2\n",
        "    nb = 1\n",
        "    # note we dont want to use a full zero encoding, so the bit_capacity will\n",
        "    # have 1 subtracted to it. that is, with 1 bit you can only encode 1 value,\n",
        "    # with 2 bits only 3 values and so on.\n",
        "    all_other_inputs = self.input_aux_size + self.extra_params\n",
        "    while bit_capacity - 1 \u003c (nb + all_other_inputs + ics) * self.size_hidden:\n",
        "      nb += 1\n",
        "      bit_capacity *= 2\n",
        "    ics += nb\n",
        "    return ics\n",
        "\n",
        "  def _createBinaryCoords(self):\n",
        "    input_coord_size = self.input_coord_size\n",
        "    layers_sizes = self.layers_sizes\n",
        "\n",
        "    def __genCoords(n):\n",
        "      nb = getNumBits(n)\n",
        "\n",
        "      all_coords = []\n",
        "      for i in range(1, n+1):\n",
        "        bit_i = f\"{i:b}\"\n",
        "        t_i = np.zeros([nb])\n",
        "        for idx, c in enumerate(bit_i[::-1]):\n",
        "          if c == '1':\n",
        "            t_i[idx] = 1.\n",
        "        all_coords.append(t_i)\n",
        "      return np.stack(all_coords).astype(np.float32)\n",
        "\n",
        "    layer_inputs = []\n",
        "    left_ics = 0\n",
        "    for i in range(len(layers_sizes) - 1):\n",
        "      sl = layers_sizes[i]\n",
        "      sr = layers_sizes[i+1]\n",
        "\n",
        "      # weight\n",
        "      current_ics = getNumBits(sl*sr)\n",
        "      layer_input =  __genCoords(sl*sr)\n",
        "      layer_input = tf.pad(layer_input, [(0,0), (left_ics, input_coord_size - left_ics - current_ics)])\n",
        "      left_ics += current_ics\n",
        "      layer_inputs.append(layer_input)\n",
        "\n",
        "      # bias\n",
        "      if i \u003c len(layers_sizes) - 2:\n",
        "        current_ics = getNumBits(sr)\n",
        "        layer_input = __genCoords(sr)\n",
        "        layer_input = tf.pad(layer_input, [(0,0), (left_ics, input_coord_size - left_ics - current_ics)])\n",
        "        layer_inputs.append(layer_input)\n",
        "    \n",
        "    return layer_inputs\n",
        "\n",
        "  def unflattenWeights(self, weights):\n",
        "    layers = []\n",
        "    for i in range((len(weights) -1)//2):\n",
        "      layers.append(\n",
        "          (weights[2*i], weights[2*i+1]))\n",
        "    layers.append((weights[-1],))\n",
        "    return layers\n",
        "\n",
        "  def set_weights(self, new_weights):\n",
        "    self.trainable_variables = [tf.Variable(v) for v in new_weights]\n",
        "    self.layers = self.unflattenWeights(self.trainable_variables)\n",
        "\n",
        "  def set_fixed_weights(self, is_fixed_logit, fixed_w):\n",
        "    self.is_fixed_logit = tf.Variable(is_fixed_logit)\n",
        "    self.fixed_w = tf.Variable(fixed_w)\n",
        "    self.trainable_fixed_variables = [self.is_fixed_logit, self.fixed_w]\n",
        "  \n",
        "  def generateNewWeights(self, variation_inputs=None, weights = None):\n",
        "    if weights is None:\n",
        "      weights = self.layers\n",
        "    weights_flat = tf.nest.flatten(weights)\n",
        "    layers_sizes = self.layers_sizes\n",
        "    new_weights = []\n",
        "    for i, coords in enumerate(self.layer_inputs):\n",
        "      l1w = self.synapses(coords, variation_inputs, weights)\n",
        "\n",
        "      l = i // 2\n",
        "\n",
        "      size_l = layers_sizes[l]\n",
        "      size_r = layers_sizes[l+1]\n",
        "\n",
        "      # compute mean of the two blocks and make sure the mean is the same.\n",
        "      if self.standardize_weights:\n",
        "        orig_m = tf.reduce_mean(weights_flat[i])\n",
        "        new_m = tf.reduce_mean(l1w)\n",
        "        l1w = l1w + (orig_m - new_m)\n",
        "\n",
        "      if i % 2 == 0:\n",
        "        # for weights, also preserve the standard deviation.\n",
        "        orig_std = tf.math.reduce_std(weights_flat[i])\n",
        "        new_std = tf.math.reduce_std(l1w)\n",
        "\n",
        "        if self.standardize_weights:\n",
        "          l1w = (l1w - orig_m) * (orig_std / new_std) + orig_m\n",
        "        l1w = tf.reshape(l1w, [size_l, size_r])\n",
        "        new_weights.append(l1w)\n",
        "      else:\n",
        "        l1b = tf.reshape(l1w, [ size_r])\n",
        "        new_weights.append(l1b)\n",
        "\n",
        "    return new_weights\n",
        "\n",
        "  def noisyCopyWeights(self, weights=None, tolerance_std=0.02):\n",
        "    if weights is None:\n",
        "      weights = self.layers\n",
        "    weights_flat = tf.nest.flatten(weights)\n",
        "    new_weights = []\n",
        "    for v in weights_flat:\n",
        "      orig_std = tf.math.reduce_std(v)\n",
        "      orig_std = tf.cond(tf.abs(orig_std) \u003c 1e-6,lambda: 1e-3,lambda: orig_std)\n",
        "      new_v = v + tf.random.normal(shape=v.shape) * orig_std * tolerance_std\n",
        "      new_weights.append(new_v)\n",
        "    return new_weights\n",
        "\n",
        "  def createNewNetwork(self, variation_inputs=None):\n",
        "    new_net = SelfReplicator(self.n_hidden, self.wo, self.size_hidden, self.standardize_weights,\n",
        "                    self.use_fixed_weights, self.switch_init_constant, self.extra_params)\n",
        "    new_weights = self.generateNewWeights(variation_inputs=variation_inputs)\n",
        "    new_net.set_weights(new_weights)\n",
        "    if self.use_fixed_weights:\n",
        "      new_net.set_fixed_weights(self.is_fixed_logit, self.fixed_w)\n",
        "    return new_net\n",
        "\n",
        "  def deserialize(self, serialized_weights):\n",
        "    layers_sizes = self.layers_sizes\n",
        "    weights_sizes = []\n",
        "    weights_dims = []\n",
        "    for i in range(len(layers_sizes)-1):\n",
        "      weights_sizes.append(layers_sizes[i]*layers_sizes[i+1])\n",
        "      weights_dims.append((layers_sizes[i], layers_sizes[i+1]))\n",
        "      if i \u003c len(layers_sizes) - 2:\n",
        "        weights_sizes.append(layers_sizes[i+1])\n",
        "        weights_dims.append((layers_sizes[i+1],))\n",
        "    all_weights_flat_split = tf.split(serialized_weights, weights_sizes)\n",
        "    all_weights_flat = [tf.reshape(t, s) for t, s in zip(\n",
        "        all_weights_flat_split, weights_dims)]\n",
        "    all_weights = []\n",
        "    for i in range((len(all_weights_flat) - 1)//2):\n",
        "      all_weights.append((all_weights_flat[2*i], all_weights_flat[2*i+1]))\n",
        "    all_weights.append((all_weights_flat[-1],))\n",
        "    return all_weights\n",
        "\n",
        "  def rgb(self, inputs, weights = None):\n",
        "    if weights is None:\n",
        "      weights = self.layers\n",
        "    inputs = tf.pad(inputs, [(0,0), (0, 1)], constant_values=1.)\n",
        "    inputs = tf.pad(inputs, [(0,0), (0, self.input_coord_size+self.extra_params)])\n",
        "    return self.call(inputs, weights)[:, :3]\n",
        "\n",
        "  def synapses(self, inputs, variation_inputs=None, weights = None):\n",
        "    if weights is None:\n",
        "      weights = self.layers\n",
        "    inputs = tf.pad(inputs, [(0,0), (self.input_aux_size, 0)])\n",
        "    # This code breaks if you call synapses without variation_inputs if it needs them.\n",
        "    if variation_inputs is not None:\n",
        "      variation_inputs = tf.broadcast_to(\n",
        "          variation_inputs, [inputs.shape[0], variation_inputs.shape[0]])\n",
        "      inputs = tf.concat([inputs, variation_inputs], -1)\n",
        "    return self.call(inputs, weights)[:, -1]\n",
        "\n",
        "  def merge_with_fixed(self, weights):\n",
        "    weights_flat = tf.nest.flatten(weights)\n",
        "    wf = tf.concat([tf.reshape(v, [-1]) for v in weights_flat], 0)\n",
        "    is_fixed = tf.sigmoid(self.is_fixed_logit)\n",
        "    new_w = wf * (1. - is_fixed) + self.fixed_w * is_fixed\n",
        "    return self.deserialize(new_w)\n",
        "\n",
        "  def call(self, inputs, weights = None):\n",
        "    if weights is None:\n",
        "      weights = self.layers\n",
        "    if self.use_fixed_weights:\n",
        "      weights = self.merge_with_fixed(weights)\n",
        "    # weights: list of (matrix, bias)\n",
        "    for n, l in enumerate(weights):\n",
        "      if n == len(weights)-1:\n",
        "        return inputs @ l[0]\n",
        "      inputs = inputs @ l[0] + l[1]\n",
        "      inputs = tf.math.sin(inputs)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZEv_0WLqpmVq"
      },
      "source": [
        "## Model from [Chang et al.](https://arxiv.org/abs/1803.05859) for benchmarking"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PqqUOcvEpk-b"
      },
      "outputs": [],
      "source": [
        "# Benchmark quine\n",
        "from tensorflow.keras.layers import Dense\n",
        "\n",
        "class BenchmarkQuine():\n",
        "  def __init__(self, n_hidden, size_hidden,\n",
        "               standardize_weights, proj_init=None):\n",
        "    self.n_hidden = n_hidden\n",
        "    self.n_outputs = 4\n",
        "    self.size_hidden = size_hidden\n",
        "    self.standardize_weights = standardize_weights\n",
        "    self.trainable_variables = []\n",
        "    self.input_coord_size = (size_hidden*size_hidden) * n_hidden + (size_hidden*self.n_outputs)\n",
        "    layers = []\n",
        "    # architecture has:\n",
        "    # - random projection layer to size_hidden\n",
        "    # n_hidden more hidden layers\n",
        "    # n_outputs\n",
        "\n",
        "    # random projections\n",
        "    if proj_init is None:\n",
        "      aux_size = 2 + 1\n",
        "      self.aux_proj = np.random.uniform(-1., 1., [aux_size, size_hidden // 2])\n",
        "      self.coord_proj = tf.keras.layers.Embedding(self.input_coord_size, size_hidden // 2)\n",
        "    else:\n",
        "      self.aux_proj = proj_init[0]\n",
        "      self.coord_proj = proj_init[1]\n",
        "\n",
        "    # weights\n",
        "    ls_except_first = [self.size_hidden]*(self.n_hidden+1) + [self.n_outputs]\n",
        "    for i in range(len(ls_except_first) -1):\n",
        "      n = ls_except_first[i]\n",
        "      n2 = ls_except_first[i+1]\n",
        "      k = math.sqrt(2/n)\n",
        "      layer_init = np.random.normal(0, k, [n, n2]).astype(np.float32)\n",
        "\n",
        "      layer_init = tf.Variable(layer_init)\n",
        "      self.trainable_variables.append(layer_init)\n",
        "      # NO BIAS.\n",
        "      layers.append((layer_init,))\n",
        "\n",
        "    self.layers = layers\n",
        "    self.layer_inputs = self._createCoords()\n",
        "    self.all_li = tf.concat(self.layer_inputs, 0)\n",
        "\n",
        "\n",
        "  def _createCoords(self):\n",
        "    all_coords = np.arange(self.input_coord_size, dtype=np.int32)\n",
        "    split_n = [self.size_hidden * self.size_hidden] * self.n_hidden + [self.size_hidden * self.n_outputs] \n",
        "    return tf.split(all_coords, split_n)\n",
        "\n",
        "  def unflattenWeights(self, weights):\n",
        "    layers = []\n",
        "    for w in weights:\n",
        "      layers.append((w,))\n",
        "    return layers\n",
        "\n",
        "  def set_weights(self, new_weights):\n",
        "    self.trainable_variables = [tf.Variable(v) for v in new_weights]\n",
        "    self.layers = self.unflattenWeights(self.trainable_variables)\n",
        "\n",
        "  def generateNewWeights(self, weights = None):\n",
        "    if weights is None:\n",
        "      weights = self.layers\n",
        "    weights_flat = tf.nest.flatten(weights)\n",
        "    layers_sizes_except_first = [self.size_hidden]*(self.n_hidden+1) + [self.n_outputs]\n",
        "    new_weights = []\n",
        "    for i, coords in enumerate(self.layer_inputs):\n",
        "      if self.n_outputs == 3:\n",
        "        r,gr, b = tf.split(self.synapses(coords, weights), 3, -1)\n",
        "        l1w = tf.reshape((r - gr) * b, [-1])\n",
        "      else:\n",
        "        l1w = self.synapses(coords, weights)\n",
        "      l = i\n",
        "\n",
        "      size_l = layers_sizes_except_first[l]\n",
        "      size_r = layers_sizes_except_first[l+1]\n",
        "\n",
        "      # compute mean of the two blocks and make sure the mean is the same.\n",
        "      if self.standardize_weights:\n",
        "        orig_m = tf.reduce_mean(weights_flat[i])\n",
        "        new_m = tf.reduce_mean(l1w)\n",
        "        l1w = l1w + (orig_m - new_m)\n",
        "\n",
        "      # also perserve the standard deviation.\n",
        "      orig_std = tf.math.reduce_std(weights_flat[i])\n",
        "      new_std = tf.math.reduce_std(l1w)\n",
        "      if self.standardize_weights:\n",
        "        l1w = (l1w - orig_m) * (orig_std / new_std) + orig_m\n",
        "      l1w = tf.reshape(l1w, [size_l, size_r])\n",
        "      new_weights.append(l1w)\n",
        "\n",
        "    return new_weights\n",
        "\n",
        "  def createNewNetwork(self):\n",
        "    new_net = BenchmarkQuine(self.n_hidden, self.size_hidden, \n",
        "                     self.standardize_weights, proj_init=(self.aux_proj, self.coord_proj))\n",
        "    new_weights = self.generateNewWeights()\n",
        "    new_net.set_weights(new_weights)\n",
        "    return new_net\n",
        "\n",
        "  def deserialize(self, serialized_weights):\n",
        "    layers_sizes_except_first = [self.size_hidden]*(self.n_hidden+1) + [self.n_outputs]\n",
        "    weights_sizes = []\n",
        "    weights_dims = []\n",
        "    for i in range(len(layers_sizes)-1):\n",
        "      weights_sizes.append(layers_sizes[i]*layers_sizes[i+1])\n",
        "      weights_dims.append((layers_sizes[i], layers_sizes[i+1]))\n",
        "    all_weights_flat_split = tf.split(serialized_weights, weights_sizes)\n",
        "    all_weights_flat = [tf.reshape(t, s) for t, s in zip(\n",
        "        all_weights_flat_split, weights_dims)]\n",
        "    all_weights = []\n",
        "    for w in all_weights:\n",
        "      all_weights.append((w,))\n",
        "    return all_weights\n",
        "\n",
        "  def rgb(self, inputs, weights = None):\n",
        "    if weights is None:\n",
        "      weights = self.layers\n",
        "    inputs = tf.pad(inputs, [(0,0), (0, 1)], constant_values=1.)\n",
        "    proj = inputs @ self.aux_proj\n",
        "    proj = tf.pad(proj, [(0,0), (0, self.size_hidden // 2)])\n",
        "    return self.call(proj, weights)[:, :3]\n",
        "\n",
        "  def synapses(self, inputs, weights = None):\n",
        "    if weights is None:\n",
        "      weights = self.layers\n",
        "    proj = self.coord_proj(inputs)\n",
        "    proj = tf.pad(proj, [(0,0), (self.size_hidden // 2, 0)])\n",
        "    return self.call(proj, weights)[:, -1]\n",
        "\n",
        "  def call(self, inputs, weights = None):\n",
        "    # This part does NOT do the random projection.\n",
        "    # call synapses or rgb for it.\n",
        "    if weights is None:\n",
        "      weights = self.layers\n",
        "    # weights: list of (matrix, bias)\n",
        "    for n, l in enumerate(weights):\n",
        "      if n == len(weights)-1:\n",
        "        return inputs @ l[0]\n",
        "      inputs = inputs @ l[0]\n",
        "      inputs = tf.nn.selu(inputs)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VaOqRMSLZ1nt"
      },
      "source": [
        "# Experiment 1: Self-replication with only variable weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TR8wlXx2KU0y"
      },
      "outputs": [],
      "source": [
        "#  Configurations used for the experiment 1.\n",
        "n_hidden = 3\n",
        "wo = 30\n",
        "size_hidden = 64\n",
        "standardize_weights=True\n",
        "use_chang_quine = False\n",
        "\n",
        "network = SelfReplicator(\n",
        "    n_hidden, wo, size_hidden, standardize_weights, \n",
        "    use_fixed_weights=False, switch_init_constant=None, num_extra_params=0)\n",
        "\n",
        "if use_chang_quine:\n",
        "  n_hidden = 2\n",
        "  size_hidden = 100\n",
        "  standardize_weights=True\n",
        "  network = BenchmarkQuine(n_hidden, size_hidden, standardize_weights)\n",
        "\n",
        "print(\"Showing parameters:\")\n",
        "from functools import reduce\n",
        "reduce_multiply = lambda l: reduce(lambda a,b: a*b, l)\n",
        "\n",
        "tot = 0\n",
        "for v in network.trainable_variables:\n",
        "  print(v.shape, reduce_multiply(v.shape))\n",
        "  tot += reduce_multiply(v.shape)\n",
        "print(\"Total parameters: \", tot)\n",
        "\n",
        "trainer = tf.keras.optimizers.Adamax(1e-3)\n",
        "\n",
        "BATCH_SIZE = 256\n",
        "\n",
        "def generate_new_dataset_iter(new_target):\n",
        "  ds = tf.data.Dataset.from_tensor_slices((im_enc_input, new_target))\n",
        "  ds = ds.cache().shuffle(1000).batch(BATCH_SIZE).repeat()\n",
        "  return iter(ds)\n",
        "\n",
        "minib_ds_iter = generate_new_dataset_iter(target)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SMbkGEExPgyB"
      },
      "outputs": [],
      "source": [
        "# hyperparameters used during training\n",
        "\n",
        "# Experiment where we use the same sink loss training regime used in Experiment 2\n",
        "sink_exp = (10., 0., 0., 10000., 0)\n",
        "# Experiment where we try to minimize the weight divergence of the child to the parent.\n",
        "wd_exp = (10., 0., 100., 0., 1)\n",
        "# Experiment where we add a loss to the child rgb image quality.\n",
        "child_loss_exp = (10., 1., 1., 0., 1)\n",
        "# Experiment where we only have weight divergence loss.\n",
        "wd_min_exp = (0., 0., 1., 0., 1)\n",
        "\n",
        "# Choose the experiment or create your own parameters.\n",
        "hpars = wd_exp\n",
        "\n",
        "pic_loss = hpars[0]\n",
        "child_loss = hpars[1] \n",
        "weight_loss = hpars[2] \n",
        "sink_loss = hpars[3] \n",
        "# This hparam is responsible for recoursively generating a child and applying\n",
        "# child_loss and weight_loss on them.\n",
        "# Sink experiments do not need it and set it to 0.\n",
        "n_iters = hpars[4]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dFq_KBSbLvEP"
      },
      "outputs": [],
      "source": [
        "# train and visualize\n",
        "\n",
        "from functools import partial\n",
        "\n",
        "def createStep(network, trainer, use_sink_loss):\n",
        "  @tf.function\n",
        "  def step(xt, yt, pic_loss, child_loss, weight_loss, sink_loss, n_iters, use_sink_loss):\n",
        "    with tf.GradientTape(persistent=True) as g:\n",
        "      y = network.rgb(xt)\n",
        "      lpic = tf.reduce_mean(tf.square(y - yt)) * pic_loss\n",
        "\n",
        "      target_wy = tf.concat([tf.reshape(v, [-1]) for v in network.trainable_variables], 0)\n",
        "\n",
        "      if use_sink_loss:\n",
        "        # sink loss: generate a neighbor and enforce it converges to the original.\n",
        "        neighbor_weights = network.unflattenWeights(network.noisyCopyWeights(tolerance_std=0.02))\n",
        "        new_w = network.generateNewWeights(weights=neighbor_weights)\n",
        "        wy = tf.concat([tf.reshape(v, [-1]) for v in new_w], 0)\n",
        "        lsink = tf.reduce_mean(tf.square(wy - target_wy)) * sink_loss\n",
        "      else:\n",
        "        lsink = 0\n",
        "\n",
        "      weights = network.layers\n",
        "      lchild_pic = 0\n",
        "      lwd = 0\n",
        "      multiplier = 1.\n",
        "      for i in range(1, n_iters+1):\n",
        "        new_w = network.generateNewWeights(weights=weights)\n",
        "        wy = tf.concat([tf.reshape(v, [-1]) for v in new_w], 0)\n",
        "        weights = network.unflattenWeights(new_w)\n",
        "\n",
        "        new_y = network.rgb(xt, weights=weights)\n",
        "        lchild_pic += tf.reduce_mean(tf.square(new_y - yt)) * child_loss * multiplier\n",
        "        lwd += tf.reduce_mean(tf.square(wy - target_wy)) * weight_loss * multiplier\n",
        "\n",
        "        multiplier /= 2.\n",
        "\n",
        "      l = lpic + lchild_pic + lwd + lsink \n",
        "\n",
        "    grads = g.gradient(l, network.trainable_variables)\n",
        "    trainer.apply_gradients(zip(grads, network.trainable_variables))\n",
        "\n",
        "    return lpic, lchild_pic, lwd, lsink\n",
        "  return partial(step, use_sink_loss=use_sink_loss)\n",
        "\n",
        "use_sink_loss = sink_loss \u003e 1e-6\n",
        "step = createStep(network, trainer, use_sink_loss=use_sink_loss)\n",
        "\n",
        "for i in range(200000):\n",
        "  xt, yt = next(minib_ds_iter)\n",
        "  lpic, lchild_pic, lwd, lsink = step(xt, yt, pic_loss, child_loss, weight_loss, sink_loss, n_iters)\n",
        "\n",
        "  if i % 5000 == 0:\n",
        "    # we need to compute the total y:\n",
        "    y = network.rgb(im_enc_input)\n",
        "    clear_output()\n",
        "    print(\"step {}, losses: rgb:{} child_rgb:{} weight_divergence:{} sink:{}\".format(i, lpic, lchild_pic, lwd, lsink))\n",
        "    pl.figure(figsize=(15, 5))\n",
        "    n_pl = 10\n",
        "    pl.subplot(1, n_pl, 1)\n",
        "    pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "    pl.axis(\"off\")\n",
        "    new_net = network\n",
        "    child_target_father = True\n",
        "    for i in range(2, n_pl + 1):\n",
        "      new_net = new_net.createNewNetwork()\n",
        "      y = new_net.rgb(im_enc_input)\n",
        "      pl.subplot(1,n_pl,i)\n",
        "      pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "      pl.axis(\"off\")\n",
        "    pl.show()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K_6LpaUAPAOA"
      },
      "source": [
        "## Evaluations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fGjOUF8ymB-6"
      },
      "outputs": [],
      "source": [
        "\n",
        "new_net = network\n",
        "pl.figure(figsize=(15, 5))\n",
        "y = new_net.rgb(im_enc_input)\n",
        "n_pl = 10\n",
        "pl.subplot(1,n_pl,1)\n",
        "pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "pl.axis(\"off\")\n",
        "for i in range(2, n_pl + 1):\n",
        "  new_net = new_net.createNewNetwork()\n",
        "  y = new_net.rgb(im_enc_input)\n",
        "  pl.subplot(1,n_pl,i)\n",
        "  pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "  pl.axis(\"off\")\n",
        "pl.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7abIvxiZj4e5"
      },
      "outputs": [],
      "source": [
        "x_l = np.arange(1, 11, dtype=np.int32)\n",
        "\n",
        "wd_fi_l = [] \n",
        "new_net = network\n",
        "\n",
        "initial_w =  tf.concat([tf.reshape(v, [-1]) for v in new_net.trainable_variables], 0)\n",
        "all_w = [initial_w]\n",
        "for _ in range(10):\n",
        "  new_net = new_net.createNewNetwork()\n",
        "\n",
        "  new_w = tf.concat([tf.reshape(v, [-1]) for v in new_net.trainable_variables], 0)\n",
        "  all_w.append(new_w)\n",
        "\n",
        "_ = show_weight_divergence(all_w)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-5S2Xc2ScfiU"
      },
      "source": [
        "# Experiment 2: Self-replication with some fixed weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3zBKkqzRcsT7"
      },
      "outputs": [],
      "source": [
        "#  Configurations used for the experiment 2.\n",
        "n_hidden = 3\n",
        "wo = 30\n",
        "size_hidden = 64\n",
        "standardize_weights=True\n",
        "switch_init_constant= 0. # 50% fixed weights at the beginning.\n",
        "\n",
        "network = SelfReplicator(\n",
        "    n_hidden, wo, size_hidden, standardize_weights, \n",
        "    use_fixed_weights=True, switch_init_constant=switch_init_constant, \n",
        "    num_extra_params=0)\n",
        "\n",
        "print(\"Showing parameters:\")\n",
        "from functools import reduce\n",
        "reduce_multiply = lambda l: reduce(lambda a,b: a*b, l)\n",
        "\n",
        "tot = 0\n",
        "for v in network.trainable_variables:\n",
        "  print(v.shape, reduce_multiply(v.shape))\n",
        "  tot += reduce_multiply(v.shape)\n",
        "print(\"Total parameters: \", tot)\n",
        "print(\"Showing fixed parameters:\")\n",
        "tot = 0\n",
        "for v in network.trainable_fixed_variables:\n",
        "  print(v.shape, reduce_multiply(v.shape))\n",
        "  tot += reduce_multiply(v.shape)\n",
        "print(\"Total parameters: \", tot)\n",
        "\n",
        "trainer = tf.keras.optimizers.Adamax(1e-3)\n",
        "\n",
        "BATCH_SIZE = 256\n",
        "\n",
        "def generate_new_dataset_iter(new_target):\n",
        "  ds = tf.data.Dataset.from_tensor_slices((im_enc_input, new_target))\n",
        "  ds = ds.cache().shuffle(1000).batch(BATCH_SIZE).repeat()\n",
        "  return iter(ds)\n",
        "\n",
        "minib_ds_iter = generate_new_dataset_iter(target)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N7KGUTBKcsUE"
      },
      "outputs": [],
      "source": [
        "# hyperparameters used during training\n",
        "\n",
        "# parameters used in Experiment 2. Note there is an extra fixed_w_loss.\n",
        "sink_exp = (10., 0., 0., 10000., 0, 1.)\n",
        "\n",
        "# Choose the experiment or create your own parameters.\n",
        "hpars = sink_exp\n",
        "\n",
        "pic_loss = hpars[0]\n",
        "child_loss = hpars[1] \n",
        "weight_loss = hpars[2] \n",
        "sink_loss = hpars[3]\n",
        "# This hparam is responsible for recoursively generating a child and applying\n",
        "# child_loss and weight_loss on them.\n",
        "# Sink experiments do not need it and set it to 0.\n",
        "n_iters = hpars[4]\n",
        "fixed_w_loss = hpars[5]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d6px94KMci5e"
      },
      "outputs": [],
      "source": [
        "# train and visualize!\n",
        "\n",
        "from functools import partial\n",
        "\n",
        "def createStep(network, trainer, use_sink_loss):\n",
        "  @tf.function\n",
        "  def step(xt, yt, pic_loss, child_loss, weight_loss, fixed_w_loss, sink_loss, n_iters, use_sink_loss):\n",
        "    with tf.GradientTape(persistent=True) as g:\n",
        "      y = network.rgb(xt)\n",
        "      lpic = tf.reduce_mean(tf.square(y - yt)) * pic_loss\n",
        "\n",
        "      target_wy = tf.concat([tf.reshape(v, [-1]) for v in network.trainable_variables], 0)\n",
        "\n",
        "      is_fixed = tf.sigmoid(network.is_fixed_logit)\n",
        "      is_fixed_mean = tf.reduce_mean(is_fixed)\n",
        "\n",
        "      if use_sink_loss:\n",
        "        # sink loss: generate a neighbor and enforce it converges to the original.\n",
        "        neighbor_weights = network.unflattenWeights(network.noisyCopyWeights(tolerance_std=0.02))\n",
        "        new_w = network.generateNewWeights(weights=neighbor_weights)\n",
        "        wy = tf.concat([tf.reshape(v, [-1]) for v in new_w], 0)\n",
        "        lsink = tf.reduce_mean(tf.square((wy - target_wy) * (1. - is_fixed))) * sink_loss\n",
        "        # normalize such that it's independent from is-fixed value:\n",
        "        lsink /= (1. - is_fixed_mean)\n",
        "      else:\n",
        "        lsink = 0\n",
        "\n",
        "\n",
        "      weights = network.layers\n",
        "      lchild_pic = 0\n",
        "      lwd = 0\n",
        "      multiplier = 1.\n",
        "      for i in range(1, n_iters+1):\n",
        "        new_w = network.generateNewWeights(weights=weights)\n",
        "        wy = tf.concat([tf.reshape(v, [-1]) for v in new_w], 0)\n",
        "        weights = network.unflattenWeights(new_w)\n",
        "\n",
        "        new_y = network.rgb(xt, weights=weights)\n",
        "        lchild_pic += tf.reduce_mean(tf.square(new_y - yt)) * child_loss * multiplier\n",
        "        # the weight loss is different, because we ignore, partially, wrong classifications\n",
        "        # for the ones that are fixed. This is likely redundant and can be removed.\n",
        "        lw =  tf.reduce_mean(tf.square((wy - target_wy) * (1. - is_fixed))) * weight_loss * multiplier\n",
        "        # normalize such that it's independent from is_fixed value:\n",
        "        lw /= (1. - is_fixed_mean)\n",
        "        lwd += lw\n",
        "\n",
        "\n",
        "        multiplier /= 2.\n",
        "\n",
        "      lfixedw = is_fixed_mean * fixed_w_loss\n",
        "\n",
        "      l = lpic + lchild_pic + lwd + lfixedw + lsink\n",
        "\n",
        "    all_vars = network.trainable_variables + network.trainable_fixed_variables\n",
        "    grads = g.gradient(l, all_vars)\n",
        "\n",
        "    trainer.apply_gradients(zip(grads, all_vars))\n",
        "\n",
        "    return lpic, lchild_pic, lwd, is_fixed_mean, lsink\n",
        "  return partial(step, use_sink_loss=use_sink_loss)\n",
        "\n",
        "use_sink_loss = sink_loss \u003e 1e-6\n",
        "step = createStep(network, trainer, use_sink_loss=use_sink_loss)\n",
        "\n",
        "for i in range(500000):\n",
        "  # using test to train too, because we want to see the effect of learning\n",
        "  # on the parameter space.\n",
        "  xt, yt = next(minib_ds_iter)\n",
        "  lpic, lchild_pic, lwd, is_fixed_mean, lsink = step(xt, yt, pic_loss, child_loss, weight_loss, fixed_w_loss, sink_loss, n_iters)\n",
        "\n",
        "  if i % 5000 == 0:\n",
        "    # we need to compute the total y:\n",
        "    y = network.rgb(im_enc_input)\n",
        "    clear_output()\n",
        "    print(\"step {}, losses: rgb:{} child_rgb:{} weight_divergence:{} is_fixed_mean:{} sink:{}\".format(i, lpic, lchild_pic, lwd, is_fixed_mean, lsink))\n",
        "    pl.figure(figsize=(15, 5))\n",
        "    n_pl = 9\n",
        "    pl.subplot(1, n_pl, 1)\n",
        "    pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "    pl.axis(\"off\")\n",
        "    new_net = network\n",
        "    child_target_father = True\n",
        "    for i in range(2, n_pl + 1):\n",
        "      new_net = new_net.createNewNetwork()\n",
        "      y = new_net.rgb(im_enc_input)\n",
        "      pl.subplot(1,n_pl,i)\n",
        "      pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "      pl.axis(\"off\")\n",
        "    pl.show()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1_5IuPQtR92Q"
      },
      "source": [
        "## Alternatively load a saved model from Github"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ww0VcuI3SBLN"
      },
      "outputs": [],
      "source": [
        "!wget -O saved_models.zip 'https://github.com/google-research/self-organising-systems/blob/master/self_replicating_nn/assets/saved_models.zip?raw=true'\n",
        "!unzip -oq \"saved_models.zip\" -d \"saved_models\"\n",
        "# Load the saved model:\n",
        "new_net = SelfReplicator(n_hidden, wo, size_hidden, standardize_weights, \n",
        "                         use_fixed_weights=True, switch_init_constant=switch_init_constant,\n",
        "                         num_extra_params=0)\n",
        "new_net.set_weights(params[:-2])\n",
        "new_net.set_fixed_weights(params[-2], params[-1])\n",
        "\n",
        "network = new_net"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QJAnFFu5Syrz"
      },
      "source": [
        "## Evaluations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "slfAFt7TS4Va"
      },
      "outputs": [],
      "source": [
        "# Show sequence\n",
        "y = network.rgb(im_enc_input)\n",
        "n_pl = 20\n",
        "sub_l = n_pl // 10 + (0 if n_pl % 10 == 0 else 1)\n",
        "fh = int(sub_l * 1.5)\n",
        "pl.figure(figsize=(15, fh))\n",
        "pl.subplot(sub_l,10,1)\n",
        "pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "pl.axis(\"off\")\n",
        "new_net = network\n",
        "\n",
        "for i in range(2, n_pl + 1):\n",
        "  new_net = new_net.createNewNetwork()\n",
        "  y = new_net.rgb(im_enc_input)\n",
        "  pl.subplot(sub_l,10,i)\n",
        "  pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "  pl.axis(\"off\")\n",
        "\n",
        "  new_w = tf.concat([tf.reshape(v, [-1]) for v in new_net.trainable_variables], 0)\n",
        "\n",
        "pl.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4hPOHqs8S_CN"
      },
      "outputs": [],
      "source": [
        "is_fixed = tf.sigmoid(network.is_fixed_logit)\n",
        "is_fixed_mean = tf.reduce_mean(is_fixed)\n",
        "print(\"Fixed mean\", is_fixed_mean)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ISET3HuTTNNA"
      },
      "outputs": [],
      "source": [
        "print(\"Showing that the variable weights matter:\")\n",
        "# Evidence that the variance of the variables is important:\n",
        "# you can also try to fill with 0, but it would be more unfair.\n",
        "zero_w = [tf.fill(v.shape, tf.reduce_mean(v)) for v in tf.nest.flatten(network.layers)]\n",
        "#zero_w = [tf.fill(v.shape, 0.) for v in tf.nest.flatten(network.layers)]\n",
        "zero_w = network.unflattenWeights(zero_w)\n",
        "\n",
        "print(\"If the self-replication were full-zero (it would then shift the output to the mean of the group):\")\n",
        "y = network.rgb(im_enc_input, weights=zero_w)\n",
        "pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "pl.axis(\"off\")\n",
        "pl.show()\n",
        "\n",
        "print(\"Showing self-replication with such initialization:\")\n",
        "n_pl = 20\n",
        "sub_l = n_pl // 10 + (0 if n_pl % 10 == 0 else 1)\n",
        "fh = int(sub_l * 1.5)\n",
        "pl.figure(figsize=(15, fh))\n",
        "pl.subplot(sub_l,10,1)\n",
        "pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "pl.axis(\"off\")\n",
        "new_w = zero_w\n",
        "\n",
        "for i in range(2, n_pl + 1):\n",
        "  new_w = network.generateNewWeights(weights=new_w)\n",
        "  y = network.rgb(im_enc_input, weights=new_w)\n",
        "  pl.subplot(sub_l,10,i)\n",
        "  pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "  pl.axis(\"off\")\n",
        "\n",
        "pl.show()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rANP1erSYddG"
      },
      "outputs": [],
      "source": [
        "# Weight divergence plot\n",
        "n_pl = 300\n",
        "sub_l = int(math.sqrt(n_pl))\n",
        "\n",
        "new_net = network\n",
        "\n",
        "initial_weights = tf.concat([tf.reshape(v, [-1]) for v in new_net.trainable_variables], 0)\n",
        "all_w = [initial_weights]\n",
        "is_fixed = tf.sigmoid(new_net.is_fixed_logit)\n",
        "is_not_fixed = 1. - is_fixed\n",
        "for i in range(1, n_pl + 1):\n",
        "  new_net = new_net.createNewNetwork()\n",
        "  y = new_net.rgb(im_enc_input)\n",
        "  new_w = tf.concat([tf.reshape(v, [-1]) for v in new_net.trainable_variables], 0)\n",
        "  all_w.append(new_w)\n",
        "\n",
        "final_w = new_w\n",
        "\n",
        "child_wd_fp, child_wd_fl = show_weight_divergence(all_w)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VqukdFA-_NEA"
      },
      "outputs": [],
      "source": [
        "# We sample neighbors and see whether they converge to the final state of a \n",
        "# normal run\n",
        "n_pl = 300\n",
        "\n",
        "num_runs = 40\n",
        "sub_l = num_runs // 10 + (0 if num_runs % 10 == 0 else 1)\n",
        "fh = int(sub_l * 1.5)\n",
        "pl.figure(figsize=(15, fh))\n",
        "\n",
        "all_runs_wd_l = []\n",
        "\n",
        "for r in range(num_runs):\n",
        "  new_net = network\n",
        "\n",
        "  wd_l = []\n",
        "\n",
        "  flat_nw = network.noisyCopyWeights(tolerance_std=0.02)\n",
        "  neighbor_weights = network.unflattenWeights(flat_nw)\n",
        "  neigh_wy = tf.concat([tf.reshape(v, [-1]) for v in flat_nw], 0)\n",
        "  lbefore = tf.reduce_mean(tf.square(neigh_wy - final_w))\n",
        "  wd_l.append(lbefore)\n",
        "\n",
        "  prev_layer = neighbor_weights\n",
        "  is_fixed = tf.sigmoid(new_net.is_fixed_logit)\n",
        "  is_not_fixed = 1. - is_fixed\n",
        "  for i in range(1, n_pl + 1):\n",
        "    new_w_l = network.generateNewWeights(weights=prev_layer)\n",
        "    new_w = tf.concat([tf.reshape(v, [-1]) for v in new_w_l], 0)\n",
        "    all_w.append(new_w)\n",
        "    weight_divergence = tf.reduce_mean(tf.square(final_w - new_w))\n",
        "    wd_l.append(weight_divergence)\n",
        "\n",
        "    prev_layer = network.unflattenWeights(new_w_l)\n",
        "  \n",
        "  y = network.rgb(im_enc_input, weights=prev_layer)\n",
        "  pl.subplot(sub_l,10,r+1)\n",
        "  pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "  pl.axis(\"off\")\n",
        "  all_runs_wd_l.append(wd_l)\n",
        "\n",
        "# Show the final images\n",
        "print(\"Final states\")\n",
        "pl.show()\n",
        "\n",
        "# Plot graph\n",
        "X = np.array(all_runs_wd_l)\n",
        "mu = X.mean(axis=0)\n",
        "sigma = X.std(axis=0)\n",
        "\n",
        "x_l = np.arange(1, len(mu)+1, dtype=np.int32)\n",
        "\n",
        "pl.title(\"Weight divergence of random neighbors from sink\")\n",
        "pl.xlabel(\"Number of replication steps\")\n",
        "pl.plot(x_l, mu)\n",
        "pl.yscale(\"log\")\n",
        "pl.fill_between(x_l, mu+sigma, mu-sigma, alpha=0.5)\n",
        "pl.show()\n",
        "\n",
        "pl.title(\"Weight divergence of random neighbors from sink\")\n",
        "pl.xlabel(\"Number of replication steps\")\n",
        "for wd_l in all_runs_wd_l:\n",
        "  x_l = np.arange(1, len(wd_l)+1, dtype=np.int32)\n",
        "  pl.plot(x_l, wd_l)\n",
        "  pl.yscale(\"log\")\n",
        "pl.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KvDD7iDTB4bI"
      },
      "outputs": [],
      "source": [
        "\n",
        "# Join the two plots\n",
        "\n",
        "pl.title(\"Weight divergence\")\n",
        "pl.plot(child_wd_fp[0], child_wd_fp[1], label=\"From parent\")\n",
        "pl.plot(child_wd_fl[0], child_wd_fl[1], label=\"From last (sink)\")\n",
        "\n",
        "X = np.array(all_runs_wd_l)\n",
        "mu = X.mean(axis=0)\n",
        "sigma = X.std(axis=0)\n",
        "\n",
        "x_l = np.arange(1, len(mu)+1, dtype=np.int32)\n",
        "\n",
        "pl.plot(x_l, mu, label=\"Random neighbor from sink\", color=\"green\")\n",
        "pl.fill_between(x_l, mu+sigma, mu-sigma,  color=\"green\", alpha=0.4)\n",
        "\n",
        "pl.yscale(\"log\")\n",
        "pl.xlabel(\"Number of replication steps\")\n",
        "pl.legend()\n",
        "pl.show()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tu2WMPbrVRtm"
      },
      "outputs": [],
      "source": [
        "is_fixed = tf.sigmoid(network.is_fixed_logit).numpy()\n",
        "pl.title(\"Histogram for the proportion of fixed weights\")\n",
        "pl.hist(is_fixed, bins=100)\n",
        "pl.show()\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n-Z8uZzHWOwN"
      },
      "outputs": [],
      "source": [
        "all_w_sizes = [tf.reshape(v, [-1]).shape[0] for v in network.trainable_variables]\n",
        "print(all_w_sizes)\n",
        "\n",
        "max_w_sw = 0.5\n",
        "print(\"Showing all variable weights (where w_sw \u003c {}) distributions.\".format(max_w_sw))\n",
        "print(\"note that the variable is multiplied by (1. - w_sw)\")\n",
        "\n",
        "for w_var, w_sw in zip(network.trainable_variables, tf.split(tf.sigmoid(network.is_fixed_logit), all_w_sizes)):\n",
        "  print(\"Var shape:\", w_var.shape)\n",
        "  variable_w = tf.gather(tf.reshape(w_var, [-1]) * (1. - w_sw), tf.where(w_sw \u003c max_w_sw)).numpy()\n",
        "  pl.hist(variable_w, bins=100)\n",
        "  pl.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r7zdoYy9FHmY"
      },
      "source": [
        "# Experiment 3: Self-replication with variation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S4JPMhVyNQDy"
      },
      "outputs": [],
      "source": [
        "# generate all possible targets.\n",
        "\n",
        "target_np = target.numpy()\n",
        "\n",
        "def right_shift_np(im):\n",
        "  return im[..., [2,0,1]]\n",
        "def left_shift_np(im):\n",
        "  return im[..., [1,2,0]]\n",
        "def swap_outer_np(im):\n",
        "  return im[..., [2,1,0]]\n",
        "\n",
        "all_variation_targets = [\n",
        "    right_shift_np(target_np),\n",
        "    left_shift_np(target_np),\n",
        "    swap_outer_np(target_np)]\n",
        "\n",
        "# prepare the variation_inputs.\n",
        "no_variation = tf.constant([0., 0., 0.])\n",
        "rs_variation = tf.constant([1., 0., 0.])\n",
        "ls_variation = tf.constant([0., 1., 0.])\n",
        "chsw_variation = tf.constant([0., 0., 1.])\n",
        "\n",
        "# does not include no_varition,\n",
        "all_variation_inputs = [rs_variation, ls_variation, chsw_variation]\n",
        "\n",
        "# plot the images\n",
        "pl.figure(figsize=(6, 12))\n",
        "pl.subplot(1,4,1)\n",
        "y = target_np\n",
        "pl.imshow(np.clip(zoom(y.reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "pl.axis(\"off\")\n",
        "\n",
        "pl.subplot(1,4,2)\n",
        "y = all_variation_targets[0]\n",
        "pl.imshow(np.clip(zoom(y.reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "pl.axis(\"off\")\n",
        "\n",
        "pl.subplot(1,4,3)\n",
        "y = all_variation_targets[1]\n",
        "pl.imshow(np.clip(zoom(y.reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "pl.axis(\"off\")\n",
        "\n",
        "pl.subplot(1,4,4)\n",
        "y = all_variation_targets[2]\n",
        "pl.imshow(np.clip(zoom(y.reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "pl.axis(\"off\")\n",
        "\n",
        "pl.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VrPZhTSMLwCA"
      },
      "outputs": [],
      "source": [
        "n_hidden = 3\n",
        "wo = 30\n",
        "size_hidden = 64\n",
        "standardize_weights=True\n",
        "switch_init_constant = 2. # Fixed weights are the majority at the beginnning.\n",
        "\n",
        "network = SelfReplicator(\n",
        "    n_hidden, wo, size_hidden, standardize_weights, \n",
        "    use_fixed_weights=True, switch_init_constant=switch_init_constant, \n",
        "    num_extra_params=3)\n",
        "\n",
        "print(\"Showing parameters:\")\n",
        "from functools import reduce\n",
        "reduce_multiply = lambda l: reduce(lambda a,b: a*b, l)\n",
        "\n",
        "tot = 0\n",
        "for v in network.trainable_variables:\n",
        "  print(v.shape, reduce_multiply(v.shape))\n",
        "  tot += reduce_multiply(v.shape)\n",
        "print(\"Total parameters: \", tot)\n",
        "print(\"Showing fixed parameters:\")\n",
        "tot = 0\n",
        "for v in network.trainable_fixed_variables:\n",
        "  print(v.shape, reduce_multiply(v.shape))\n",
        "  tot += reduce_multiply(v.shape)\n",
        "print(\"Total parameters: \", tot)\n",
        "\n",
        "lr_sched = tf.keras.optimizers.schedules.PiecewiseConstantDecay(\n",
        "  [100000], [1e-3, 1e-4])\n",
        "trainer = tf.keras.optimizers.Adamax(lr_sched)\n",
        "\n",
        "BATCH_SIZE = 256\n",
        "\n",
        "def generate_new_dataset_iter(new_target):\n",
        "  ds = tf.data.Dataset.from_tensor_slices((im_enc_input, new_target, tf.constant(np.arange(len(im_enc_input), dtype=np.int32)[..., np.newaxis])))\n",
        "  ds = ds.cache().shuffle(1000).batch(BATCH_SIZE).repeat()\n",
        "  return iter(ds)\n",
        "\n",
        "minib_ds_iter = generate_new_dataset_iter(target)\n",
        "\n",
        "tot_steps = 0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uPqRWvv8R7s5"
      },
      "outputs": [],
      "source": [
        "# train and visualize.\n",
        "import random\n",
        "\n",
        "from functools import partial\n",
        "\n",
        "@tf.function\n",
        "def step(xt, yt, idxt, variation_inputs, t1, pic_loss, fixed_w_loss, sink_loss, switch_colors_loss, n_iters):\n",
        "  with tf.GradientTape(persistent=True) as g:\n",
        "    y = network.rgb(xt)\n",
        "    lpic = tf.reduce_mean(tf.square(y - yt)) * pic_loss\n",
        "\n",
        "    target_wy = tf.concat([tf.reshape(v, [-1]) for v in network.trainable_variables], 0)\n",
        "\n",
        "    is_fixed = tf.sigmoid(network.is_fixed_logit)\n",
        "    is_fixed_mean = tf.reduce_mean(is_fixed)\n",
        "\n",
        "    # sink loss: generate a neighbor and enforce it converges to the original.\n",
        "    neighbor_weights = network.unflattenWeights(network.noisyCopyWeights(tolerance_std=0.02))\n",
        "    new_w = network.generateNewWeights(no_variation, weights=neighbor_weights)\n",
        "    wy = tf.concat([tf.reshape(v, [-1]) for v in new_w], 0)\n",
        "    lsink = tf.reduce_mean(tf.square((wy - target_wy) * (1. - is_fixed))) * sink_loss\n",
        "    # normalize such that it's independent from is-fixed value:\n",
        "    lsink /= (1. - is_fixed_mean)\n",
        "\n",
        "    # Now apply, twice, a random transformation:\n",
        "    # - right shift ch image\n",
        "    # - left shift ch image\n",
        "    # - swap outer channels\n",
        "    # And then do a sink loss on that.\n",
        "\n",
        "    weights = network.layers\n",
        "\n",
        "    new_w = network.generateNewWeights(variation_inputs, weights=weights)\n",
        "    wy = tf.concat([tf.reshape(v, [-1]) for v in new_w], 0)\n",
        "    weights = network.unflattenWeights(new_w)\n",
        "    new_y = network.rgb(xt, weights=weights)\n",
        "\n",
        "    # generate the target y now:\n",
        "    new_target = tf.gather(t1, idxt)\n",
        "    lchild_pic = tf.reduce_mean(tf.square(new_y - new_target)) * switch_colors_loss\n",
        "    lchild_pic *= 3 # there are 3 options here.\n",
        "\n",
        "    # Now sink loss on child.\n",
        "    neighbor_weights = network.unflattenWeights(network.noisyCopyWeights(weights=weights, tolerance_std=0.02))\n",
        "    new_new_w = network.generateNewWeights(no_variation, weights=neighbor_weights)\n",
        "    wy = tf.concat([tf.reshape(v, [-1]) for v in new_new_w], 0)\n",
        "    target_wy = tf.concat([tf.reshape(v, [-1]) for v in new_w], 0)\n",
        "    lsink_child = tf.reduce_mean(tf.square((wy - target_wy) * (1. - is_fixed))) * sink_loss\n",
        "    # normalize such that it's independent from is-fixed value:\n",
        "    lsink_child /= (1. - is_fixed_mean)\n",
        "    lsink_child *= 3 # there are 3 options here.\n",
        "\n",
        "    lfixedw = is_fixed_mean * fixed_w_loss\n",
        "\n",
        "    l = lpic + lfixedw + lsink + lsink_child + lchild_pic\n",
        "\n",
        "  all_vars = network.trainable_variables + network.trainable_fixed_variables\n",
        "  grads = g.gradient(l, all_vars)\n",
        "\n",
        "  trainer.apply_gradients(zip(grads, all_vars))\n",
        "\n",
        "  return lpic, is_fixed_mean, lsink, lsink_child, lchild_pic\n",
        "\n",
        "for i in range(500000):\n",
        "  tot_steps += 1\n",
        "  # using test to train too, because we want to see the effect of learning\n",
        "  # on the parameter space.\n",
        "  xt, yt, idxt = next(minib_ds_iter)\n",
        "  idxt = tf.reshape(idxt, [-1])\n",
        "\n",
        "  # generate the transformations.\n",
        "  r = random.random()\n",
        "  if r \u003c 1/3:\n",
        "    id1 = 0\n",
        "  elif r \u003c 2/3:\n",
        "    id1 = 1\n",
        "  else:\n",
        "    id1 = 2\n",
        "  variation_inputs = all_variation_inputs[id1]\n",
        "  t1 = all_variation_targets[id1]\n",
        " \n",
        "  pic_loss = 10.\n",
        "  fixed_w_loss = 1.\n",
        "  sink_loss = 10000 # 10000.\n",
        "  switch_colors_loss = 10. # All losses are also multiplied by the chance of it happening.\n",
        "  n_iters = 0\n",
        "  lpic, is_fixed_mean, lsink, lsink_child, lchild_pic = step(\n",
        "      xt, yt, idxt, variation_inputs, t1, pic_loss, fixed_w_loss, sink_loss, switch_colors_loss, n_iters)\n",
        "\n",
        "  if i % 5000 == 0:\n",
        "    #clear_output()\n",
        "    # we need to compute the total y:\n",
        "    clear_output()\n",
        "    print(\"step {} (tot: {}), losses: rgb:{} is_fixed_mean:{} sink_loss:{}\".format(i, tot_steps, lpic, is_fixed_mean, lsink))\n",
        "    print(\"Child 1: rgb : {} (norm {}) sink_child_loss:{} (norm {})\".format(lchild_pic, lchild_pic / 3 / switch_colors_loss, lsink_child, lsink_child / 3))\n",
        "    images = []\n",
        "    #show_batch(image_batch.numpy())\n",
        "    print(\"Original generate 0 degrees offspring.\")\n",
        "    pl.figure(figsize=(15, 15))\n",
        "    n_pl = 9\n",
        "    pl.subplot(1,n_pl,1)\n",
        "    y = network.rgb(im_enc_input)\n",
        "    pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "    pl.axis(\"off\")\n",
        "    new_net = network\n",
        "    child_target_father = True\n",
        "    for i in range(2, n_pl + 1):\n",
        "      new_net = new_net.createNewNetwork(no_variation)\n",
        "   \n",
        "      y = new_net.rgb(im_enc_input)\n",
        "      pl.subplot(1,n_pl,i)\n",
        "      pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "      pl.axis(\"off\")\n",
        "    pl.show()\n",
        "\n",
        "    print(\"colored child: generate new offspring.\")\n",
        "    pl.figure(figsize=(15, 5))\n",
        "    new_net = network.createNewNetwork(rs_variation)\n",
        "    y = new_net.rgb(im_enc_input)\n",
        "    pl.subplot(1,n_pl,1)\n",
        "    pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "    pl.axis(\"off\")\n",
        "    for i in range(2, n_pl + 1):\n",
        "      new_net = new_net.createNewNetwork(no_variation)\n",
        "\n",
        "      y = new_net.rgb(im_enc_input)\n",
        "      pl.subplot(1,n_pl,i)\n",
        "      pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "      pl.axis(\"off\")\n",
        "    pl.show()\n",
        "    print(\"colored child: generate new offspring.\")\n",
        "    pl.figure(figsize=(15, 5))\n",
        "    new_net = network.createNewNetwork(ls_variation)\n",
        "    y = new_net.rgb(im_enc_input)\n",
        "    pl.subplot(1,n_pl,1)\n",
        "    pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "    pl.axis(\"off\")\n",
        "    for i in range(2, n_pl + 1):\n",
        "      new_net = new_net.createNewNetwork(no_variation)\n",
        "\n",
        "      y = new_net.rgb(im_enc_input)\n",
        "      pl.subplot(1,n_pl,i)\n",
        "      pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "      pl.axis(\"off\")\n",
        "    pl.show()\n",
        "    print(\"colored child: generate new offspring.\")\n",
        "    pl.figure(figsize=(15, 5))\n",
        "    new_net = network.createNewNetwork(chsw_variation)\n",
        "    y = new_net.rgb(im_enc_input)\n",
        "    pl.subplot(1,n_pl,1)\n",
        "    pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "    pl.axis(\"off\")\n",
        "    for i in range(2, n_pl + 1):\n",
        "      new_net = new_net.createNewNetwork(no_variation)\n",
        "\n",
        "      y = new_net.rgb(im_enc_input)\n",
        "      pl.subplot(1,n_pl,i)\n",
        "      pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "      pl.axis(\"off\")\n",
        "    pl.show()\n",
        "\n",
        "    print(\"2nd gen colored child: rs rs\")\n",
        "    pl.figure(figsize=(15, 5))\n",
        "    new_net = network.createNewNetwork(rs_variation)\n",
        "    y = new_net.rgb(im_enc_input)\n",
        "    pl.subplot(1,n_pl,1)\n",
        "    pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "    pl.axis(\"off\")\n",
        "    new_net = new_net.createNewNetwork(rs_variation)\n",
        "    y = new_net.rgb(im_enc_input)\n",
        "    pl.subplot(1,n_pl,2)\n",
        "    pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "    pl.axis(\"off\")\n",
        "    for i in range(3, n_pl + 1):\n",
        "      new_net = new_net.createNewNetwork(no_variation)\n",
        "\n",
        "      y = new_net.rgb(im_enc_input)\n",
        "      pl.subplot(1,n_pl,i)\n",
        "      pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "      pl.axis(\"off\")\n",
        "    pl.show()\n",
        "\n",
        "    print(\"2nd gen colored child: rs ls\")\n",
        "    pl.figure(figsize=(15, 5))\n",
        "    new_net = network.createNewNetwork(rs_variation)\n",
        "    y = new_net.rgb(im_enc_input)\n",
        "    pl.subplot(1,n_pl,1)\n",
        "    pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "    pl.axis(\"off\")\n",
        "    new_net = new_net.createNewNetwork(ls_variation)\n",
        "    y = new_net.rgb(im_enc_input)\n",
        "    pl.subplot(1,n_pl,2)\n",
        "    pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "    pl.axis(\"off\")\n",
        "    for i in range(3, n_pl + 1):\n",
        "      new_net = new_net.createNewNetwork(no_variation)\n",
        "\n",
        "      y = new_net.rgb(im_enc_input)\n",
        "      pl.subplot(1,n_pl,i)\n",
        "      pl.imshow(np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1))\n",
        "      pl.axis(\"off\")\n",
        "    pl.show()\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GeHOadZKcMrM"
      },
      "source": [
        "## Alternatively load a saved model from Github"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "38MBj7bQMnQd"
      },
      "outputs": [],
      "source": [
        "# Load the saved model:\n",
        "!wget -O saved_models.zip 'https://github.com/google-research/self-organising-systems/blob/master/self_replicating_nn/assets/saved_models.zip?raw=true'\n",
        "!unzip -oq \"saved_models.zip\" -d \"saved_models\"\n",
        "new_net = SelfReplicator(n_hidden, wo, size_hidden, standardize_weights, \n",
        "                use_fixed_weights=True, switch_init_constant=switch_init_constant, num_extra_params=3)\n",
        "new_net.set_weights(params[:-2])\n",
        "new_net.set_fixed_weights(params[-2], params[-1])\n",
        "\n",
        "network = new_net"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D7L8J_ANcPGk"
      },
      "source": [
        "## Evaluations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_XQ24TiTckA7"
      },
      "outputs": [],
      "source": [
        "def get_ops(ops_to_do, idx, default=no_variation):\n",
        "  if idx \u003c len(ops_to_do):\n",
        "    return ops_to_do[idx]\n",
        "  return default\n",
        "\n",
        "def convert_y_to_im(y):\n",
        "  return np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1)\n",
        "\n",
        "def viz_sequence(all_im):\n",
        "  n_pl = len(all_im)\n",
        "  n_cols = 10\n",
        "  n_rows = int(math.ceil(n_pl / n_cols))\n",
        "  fh = int(n_rows * 1.5)\n",
        "  pl.figure(figsize=(15, fh))\n",
        "  for i, im in enumerate(all_im):\n",
        "    pl.subplot(n_rows,n_cols,i+1)\n",
        "    pl.imshow(im)\n",
        "    pl.axis(\"off\")\n",
        "\n",
        "  pl.show()\n",
        "\n",
        "def show_sequence(ops_to_do, default, n_pl, network=network, return_network=False):\n",
        "  new_net = network\n",
        "  all_im = []\n",
        "  for i in range(1, n_pl + 1):\n",
        "    y = new_net.rgb(im_enc_input)\n",
        "    im = convert_y_to_im(y)\n",
        "    all_im.append(im)\n",
        "\n",
        "    new_net = new_net.createNewNetwork(get_ops(ops_to_do, i-1, default=default))\n",
        "\n",
        "  viz_sequence(all_im)\n",
        "  \n",
        "  if return_network:\n",
        "    return all_im, new_net\n",
        "  return all_im\n",
        "\n",
        "def evaluate_identity(n_pl, network=network, variation=no_variation, verbose=False):\n",
        "  new_net = network\n",
        "\n",
        "  all_im = []\n",
        "  y = new_net.rgb(im_enc_input)\n",
        "  im = convert_y_to_im(y)\n",
        "  if not verbose:\n",
        "    print(\"Initial image\")\n",
        "    pl.imshow(im)\n",
        "    pl.axis(\"off\")\n",
        "    pl.show()\n",
        "  else:\n",
        "    all_im.append(im)\n",
        "\n",
        "  initial_weights = tf.concat([tf.reshape(v, [-1]) for v in new_net.trainable_variables], 0)\n",
        "  prev_w = initial_weights\n",
        "  is_fixed = tf.sigmoid(new_net.is_fixed_logit)\n",
        "  is_not_fixed = 1. - is_fixed\n",
        "  all_w = [initial_weights]\n",
        "  for i in range(1, n_pl + 1):\n",
        "    new_net = new_net.createNewNetwork(variation)\n",
        "    new_w = tf.concat([tf.reshape(v, [-1]) for v in new_net.trainable_variables], 0)\n",
        "    all_w.append(new_w)\n",
        "\n",
        "    prev_w = new_w\n",
        "\n",
        "    if verbose:\n",
        "      y = new_net.rgb(im_enc_input)\n",
        "      im = convert_y_to_im(y)\n",
        "      all_im.append(im)\n",
        "\n",
        "  show_weight_divergence(all_w)\n",
        "\n",
        "  if not verbose:\n",
        "    print(\"Final image\")\n",
        "    y = new_net.rgb(im_enc_input)\n",
        "    im = np.clip(zoom(y.numpy().reshape([IMG_HEIGHT, IMG_WIDTH, 3])), 0,1)\n",
        "    pl.imshow(im)\n",
        "    pl.axis(\"off\")\n",
        "    pl.show()\n",
        "  else:\n",
        "    print(\"Sequence:\")\n",
        "    viz_sequence(all_im)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-7oEPVTtfBfR"
      },
      "outputs": [],
      "source": [
        "# Examples\n",
        "ops_to_do = [rs_variation, rs_variation]\n",
        "\n",
        "all_im = show_sequence(ops_to_do, no_variation, 40)\n",
        "show_single_im(all_im[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_rpGWyRZeEgv"
      },
      "outputs": [],
      "source": [
        "ops_to_do = []\n",
        "\n",
        "all_im = show_sequence(ops_to_do, chsw_variation, 40)\n",
        "show_single_im(all_im[-1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g3dCYmjHd6Hp"
      },
      "source": [
        "These evaluations investigate the results of sink loss applied on our 4 in-training configurations.\n",
        "\n",
        "We show the sequences of when we repeatedly apply a \"no variation\" mutation from a state trained with sink loss.\n",
        "\n",
        "Note that:\n",
        "\n",
        "1.   The network where we apply the sink loss during training is *NEVER* the sink. We find the sink after repeated applications of the replication function.\n",
        "2.   We infer 1. from the plots, showing that there is a high divergence from the original network, but the divergence from the parent eventually converges to insignificant values.\n",
        "3.   We also see that the \"channel swap\" network did not create a sink and it diverts to another sink (the left shift sink).\n",
        "4.   We infer 3. from the visual change in the picture and a clear 2-mode behavior observed in the divergence from parent graph. Further evidence of the convergence to the \"left shift sink\" can be seen in other blocks down below.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pre3KIxQrQaL"
      },
      "outputs": [],
      "source": [
        "# Showing if the states with sink loss are actual sinks\n",
        "\n",
        "# If you set verbose to true, we will print the whole sequence of images.\n",
        "verbose = True\n",
        "n_pl_after_first = 299\n",
        "\n",
        "print(\"net: id*{}\".format(n_pl_after_first+1))\n",
        "evaluate_identity(n_pl=1+n_pl_after_first, verbose=verbose)\n",
        "\n",
        "print(\"net: ls id*{}\".format(n_pl_after_first))\n",
        "new_net = network.createNewNetwork(ls_variation)\n",
        "evaluate_identity(n_pl=n_pl_after_first, network=new_net, verbose=verbose)\n",
        "\n",
        "print(\"net: rs id*{}\".format(n_pl_after_first))\n",
        "new_net = network.createNewNetwork(rs_variation)\n",
        "evaluate_identity(n_pl=n_pl_after_first, network=new_net, verbose=verbose)\n",
        "\n",
        "# This one IS NOT a sink!\n",
        "print(\"net: chsw id*{}\".format(n_pl_after_first))\n",
        "new_net = network.createNewNetwork(chsw_variation)\n",
        "evaluate_identity(n_pl=n_pl_after_first, network=new_net, verbose=verbose)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7FPmIIbyfeaU"
      },
      "source": [
        "Below, we investigate whether repeated applications of replication with variation converge to a sink.\n",
        "\n",
        "We observe:\n",
        "\n",
        "1.   All these replication with variation operations show a clear sink.\n",
        "2.   This is undesired behavior and shows overfitting to training configurations.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mSWkV3aSTnD9"
      },
      "outputs": [],
      "source": [
        "# Showing if repeated applications of a variation cause a sink\n",
        "\n",
        "# If you set this to true, we will print the whole sequence of images.\n",
        "verbose = True\n",
        "n_pl = 299\n",
        "\n",
        "print(\"net: ls*{}\".format(n_pl))\n",
        "new_net = network.createNewNetwork(ls_variation)\n",
        "evaluate_identity(n_pl=n_pl, variation=ls_variation, verbose=verbose)\n",
        "\n",
        "print(\"net: rs*{}\".format(n_pl))\n",
        "new_net = network.createNewNetwork(rs_variation)\n",
        "evaluate_identity(n_pl=n_pl, variation=rs_variation, verbose=verbose)\n",
        "\n",
        "print(\"net: chsw*{}\".format(n_pl))\n",
        "new_net = network.createNewNetwork(chsw_variation)\n",
        "evaluate_identity(n_pl=n_pl, variation=chsw_variation, verbose=verbose)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hKgxC0igf-9m"
      },
      "source": [
        "Below, we analyze the sinks found in these out of training transitions.\n",
        "\n",
        "We can conclude that all these sinks are equivalent from observing the negligible values of their weight divergences.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1VI8vh86eQlx"
      },
      "outputs": [],
      "source": [
        "# Analising the sinks to prove they are the same:\n",
        "\n",
        "print(\"net_1: ls ls id*98\")\n",
        "ops_to_do = [ls_variation] * 2\n",
        "all_im1, net_1 = show_sequence(ops_to_do, default=no_variation, n_pl=100, return_network=True)\n",
        "\n",
        "print(\"net_2: rs*5 id*95\")\n",
        "ops_to_do = [rs_variation] * 5\n",
        "all_im2, net_2 = show_sequence(ops_to_do, default=no_variation, n_pl=100, return_network=True)\n",
        "\n",
        "\n",
        "net1_w = tf.concat([tf.reshape(v, [-1]) for v in net_1.trainable_variables], 0)\n",
        "net2_w = tf.concat([tf.reshape(v, [-1]) for v in net_2.trainable_variables], 0)\n",
        "\n",
        "weight_divergence = tf.reduce_mean(tf.square((net2_w - net1_w)))\n",
        "\n",
        "print(\"weight divergence of net_1 vs net_2:\", weight_divergence)\n",
        "\n",
        "print(\"trained sink ls_net: ls id*99\")\n",
        "ops_to_do = [ls_variation]\n",
        "all_im3, ls_net = show_sequence(ops_to_do, default=no_variation, n_pl=100, return_network=True)\n",
        "ls_net_w = tf.concat([tf.reshape(v, [-1]) for v in ls_net.trainable_variables], 0)\n",
        "weight_divergence = tf.reduce_mean(tf.square((ls_net_w - net1_w)))\n",
        "print(\"weight divergence of net_1 vs ls_net:\", weight_divergence)\n",
        "weight_divergence = tf.reduce_mean(tf.square((ls_net_w - net2_w)))\n",
        "print(\"weight divergence of net_2 vs ls_net:\", weight_divergence)\n",
        "\n",
        "print(\"net_3: chsw*9 id*91\")\n",
        "ops_to_do = [chsw_variation] * 9\n",
        "\n",
        "all_im, net_3 = show_sequence(ops_to_do, default=no_variation, n_pl=100, return_network=True)\n",
        "net3_w = tf.concat([tf.reshape(v, [-1]) for v in net_3.trainable_variables], 0)\n",
        "\n",
        "weight_divergence = tf.reduce_mean(tf.square((ls_net_w - net3_w)))\n",
        "print(\"weight divergence of net_3 vs ls_net:\", weight_divergence)\n",
        "weight_divergence = tf.reduce_mean(tf.square((net3_w - net2_w)))\n",
        "print(\"weight divergence of net_3 vs net_2:\", weight_divergence)\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "machine_shape": "hm",
      "name": "recursively_fertile_self_replicating.ipynb",
      "provenance": [
        {
          "file_id": "1glSUj0DDT8gGxU4ZFspwwz48FcUoig3y",
          "timestamp": 1620806912096
        }
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
